use std::fmt::{Display, Formatter};
use std::{i32};
use std::collections::HashSet;
use std::net::{Ipv4Addr, SocketAddrV4};
use std::str::FromStr;
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use crate::torrent::Torrent;

const BLOCK_SIZE: i32 = 16 * 1024;
#[derive(Debug)]
pub struct Peer {
    pub addr: SocketAddrV4,
    pub id: Option<[u8; 20]>,
    pub session: Option<TcpStream>,
}

impl From<&[u8]> for Peer {
    fn from(bytes: &[u8]) -> Peer {
        let ip = Ipv4Addr::new(bytes[0], bytes[1], bytes[2], bytes[3]);
        let port = u16::from_be_bytes([bytes[4], bytes[5]]);
        let addr = SocketAddrV4::from_str(format!("{}:{}", ip, port).as_str()).unwrap();
        Peer { addr, id: None, session: None }
    }
}

impl Display for Peer {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.addr.to_string().as_str())
    }
}

impl Peer {
    pub async fn handshake(&mut self, info_hash: [u8; 20]) -> anyhow::Result<()> {
        let mut stream = TcpStream::connect(self.addr).await?;
        println!("Connected");
        let mut handshake = HandShake::new(info_hash);
        let bytes = handshake.as_bytes();
        stream.write_all(bytes).await?;
        stream.read_exact(bytes).await?;
        println!("Peer ID: {}", hex::encode(&handshake.peer_id));
        self.session = Some(stream);
        Ok(())
    }

    async fn send(&mut self, message: Message) -> anyhow::Result<()> {
        if let Some(session) = self.session.as_mut() {
            session.write_all(message.to_bytes().as_slice()).await?;
            Ok(())
        } else {
            Err(anyhow::format_err!("Connection to {} is not established", self.addr))
        }
    }

    async fn recv(&mut self) -> anyhow::Result<Message> {
        if let Some(session) = self.session.as_mut() {
            // Parsing prefix
            let mut buf = [0u8; 4];
            match session.read_exact(&mut buf).await {
                Ok(..) => {}
                Err(err) => println!("ERROR while receiving message {err}")
            }
            let prefix: i32 = i32::from_be_bytes(buf);
            // Parsing message id
            let mut buf = [0u8; 1];
            session.read_exact(&mut buf).await?;
            let id: MESSAGE = unsafe { std::mem::transmute(buf[0]) };
            // Parsing payload
            let mut buf = vec![0u8; prefix as usize - 1];
            session.read_exact(&mut buf).await?;
            Ok(Message {prefix, id, payload: buf})
        } else {
            Err(anyhow::format_err!("Connection to {} is not established", self.addr))
        }
    }

    /* Something like 10110100 (size = X), where 1 - index (highest bit in byte) means
   that client has this piece
   bitfield: <len=0001+X><id=5><bitfield>
   The bitfield message is variable length, where X is the length of the bitfield
   */
    pub async fn get_pieces(&mut self) -> anyhow::Result<HashSet<i32>> {
        let bitfield = self.recv().await?;
        anyhow::ensure!(bitfield.id == MESSAGE::BITFIELD);
        let bitfield = bitfield.payload;
        let bitset: Vec<char> = bitfield.into_iter().
            flat_map(|byte| format!("{:b}", byte).chars().collect::<Vec<char>>()).collect();
        let pieces = bitset.into_iter().enumerate()
            .filter(|(_, b)| b == &'1').map(|(i, _)| i as i32).collect::<HashSet<i32>>();
        Ok(pieces)
    }

    async fn load_block(&mut self, index: i32, begin: i32, length: i32) -> anyhow::Result<Message> {
        // request: <len=0013><id=6><index><begin><length>
        let payload: Vec<u8> = [index.to_be_bytes().as_slice(), begin.to_be_bytes().as_slice(), length.to_be_bytes().as_slice()].concat();
        let request = Message::new(MESSAGE::REQUEST, payload);
        self.send(request).await?;
        let block = self.recv().await?;
        anyhow::ensure!(block.id == MESSAGE::PIECE);
        Ok(block)
    }

    pub async fn send_interested(&mut self) -> anyhow::Result<()> {
        self.send(Message::new(MESSAGE::INTERESTED, vec![])).await?;
        Ok(())
    }

    pub async fn get_unchoke(&mut self) -> anyhow::Result<()> {
        let unchoke = self.recv().await?;
        anyhow::ensure!(unchoke.id == MESSAGE::UNCHOKE);
        Ok(())
    }

    pub async fn load_piece(&mut self, index: i32, torrent: &Torrent) -> anyhow::Result<Vec<u8>> {
        // Get size of the piece
        let file_size = torrent.meta.info.length.clone() as i32;
        let piece_len: i32 = torrent.meta.info.piece_length.clone() as i32;
        let piece_size = piece_len.min(file_size - piece_len * index as i32);
        // Break the piece into blocks of 16 kiB (16 * 1024 bytes)
        // and send a request message for each block
        let mut piece: Vec<u8> = Vec::new();
        let mut remain = piece_size;
        let mut offset = 0;
        loop {
            if remain == 0 {
                // all blocks have been loaded
                break;
            }
            let block_size = BLOCK_SIZE.min(remain);
            let response = self.load_block(index, offset, block_size).await?;
            remain -= block_size;
            offset += block_size;
            let block = &response.payload[8..];
            anyhow::ensure!(response.prefix as usize - 9 == block.len());
            piece.append(&mut block.to_vec());
        }
        Ok(piece)
    }
}

// The handshake is a required message and must be the first message transmitted by the client.
#[derive(Debug, Deserialize, Serialize)]
#[allow(dead_code)]
#[repr(C)]
pub struct HandShake {
    // string length of <pstr>, as a single raw byte
    pstrlen: [u8; 1],
    // string identifier of the protocol
    pstr: [u8; 19],
    // eight (8) reserved bytes. All current implementations use all zeroes
    reserved: [u8; 8],
    // 20-byte SHA1 hash of the info key in the metainfo file
    info_hash: [u8; 20],
    // 20-byte string used as a unique ID for the client.
    peer_id: [u8; 20],
}

impl HandShake {
    pub fn new(hash: [u8; 20]) -> Self {
        let proto_len: [u8; 1] = [19];
        let bit_torrent_str: [u8; 19] = *b"BitTorrent protocol";
        let reserved = [0; 8];
        let sha1_info_hash = hash;
        let peer_id= *b"00112233445566778899";
        HandShake { pstrlen: proto_len, pstr: bit_torrent_str, reserved, info_hash: sha1_info_hash, peer_id }
    }
    pub fn as_bytes(&mut self) -> &mut [u8] {
        let bytes = self as *mut Self as *mut [u8; std::mem::size_of::<Self>()];
        // Safety: Self is a POD with repr(c) and repr(packed)
        let bytes: &mut [u8; std::mem::size_of::<Self>()] = unsafe { &mut *bytes };
        bytes
    }
}

impl From<[u8; 68]> for HandShake {
    fn from(value: [u8; 68]) -> Self {
        let mut hand_shake = HandShake::new([0u8; 20]);
        hand_shake.peer_id.clone_from_slice(&value[48..]); // get peer id (last 20 bytes) from response
        hand_shake
    }
}

// All the remaining messages in the protocol take the form of <length prefix><message ID><payload>
#[derive(Clone, Debug)]
struct Message {
    // The length prefix is a four byte big-endian value
    pub prefix: i32,
    // The message ID is a single decimal byte
    pub id: MESSAGE,
    // The payload is message dependent.
    pub payload: Vec<u8>
}

impl Message {
    pub fn new(id: MESSAGE, payload: Vec<u8>) -> Self {
        Self {prefix: payload.len() as i32 + 1, id, payload}
    }
    pub fn to_bytes(&self) -> Vec<u8> {
        [
            self.prefix.to_be_bytes().as_slice(),
            [self.id.clone() as u8; 1].as_slice(),
            self.payload.as_slice()
        ].concat()
    }
}

#[derive(PartialEq, Clone, Debug)]
#[repr(u8)]
enum MESSAGE {
    BITFIELD = 5,
    INTERESTED = 2,
    UNCHOKE = 1,
    REQUEST = 6,
    PIECE = 7
}
